from __future__ import division
from functools import partial

from aima.search import Problem
from osm_utils import CountryMap, Junction, Link
from compass_bearing import calc_turn_angle

def C_c(problem, node1, link, junction2_key):
    return link.distance

def H_c(problem, node):
    return problem.map.JunctionDistance(node.state, problem.goal)

def C_f(problem, node1, link, junction2_key):
    return link.CalcTime()

def H_f(problem, node):
    # cache the max speed result for the current problem run. we cache it here and not 
    #  inside the map because we assume it could change after load, but not during problem search
    if not hasattr(problem, '_H_f_max_speed'):
        problem._H_f_max_speed = problem.map.GetMaxSpeed()
    return problem.map.JunctionDistance(node.state, problem.goal) / problem._H_f_max_speed

def C_s(problem, node1, link, junction2_key):
    junction1 = problem.map.GetJunction(node1.state)
    junction2 = problem.map.GetJunction(junction2_key)

    if junction1 == junction2:
        return 0

    prev_unique_node = node1.parent
    while prev_unique_node:
        prev_unique_junction = problem.map.GetJunction(prev_unique_node.state)
        if prev_unique_junction != junction1:
            break
        prev_unique_node = prev_unique_node.parent

    if prev_unique_node is None:
        return 0

    angle = calc_turn_angle(prev_unique_junction, junction1, junction2) 
    
    return 1 if angle >= MapProblem.TURN_ANGLE else 0

def H_s(problem, node):
    return 0

def C_sigma(W_c, W_f, problem, node1, link, junction2_key):
    assert W_c <= 1 and W_f <= 1# and W_s <= 1 and W_c + W_f + W_s == 1
    W_s = 1 - W_c - W_f

    set_problem_sigma_gammas(problem)

    return (W_c * problem._sigma_gamma_c * C_c(problem, node1, link, junction2_key)
            + W_f * problem._sigma_gamma_f * C_f(problem, node1, link, junction2_key)
            + W_s * problem._sigma_gamma_s * C_s(problem, node1, link, junction2_key))

def H_sigma(W_c, W_f, problem, node):
    assert W_c <= 1 and W_f <= 1# and W_s <= 1 and W_c + W_f + W_s == 1
    W_s = 1 - W_c - W_f

    set_problem_sigma_gammas(problem)
    
    return (W_c * problem._sigma_gamma_c * H_c(problem, node)
            + W_f * problem._sigma_gamma_f * H_f(problem, node))

def set_problem_sigma_gammas(problem):
    if not hasattr(problem, '_sigma_gamma_c'):
        problem._sigma_gamma_c = calc_gamma_c(problem.map)
        #print problem._sigma_gamma_c
    if not hasattr(problem, '_sigma_gamma_f'):
        problem._sigma_gamma_f = calc_gamma_f(problem.map)
        #print problem._sigma_gamma_f
    if not hasattr(problem, '_sigma_gamma_s'):
        problem._sigma_gamma_s = calc_gamma_s(problem.map)
        #print problem._sigma_gamma_s
        

def calc_gamma_c(map):
    return 1.0 / map.GetAvgDistance()

def calc_gamma_f(map):
    # distance is in meters, speed is in km/h. we want time in seconds, so we multiply by 3.6
    avg_time = float(map.GetAvgDistance()) / map.GetAvgSpeed() * 3.6
    return 1.0 / avg_time

def calc_gamma_s(map):
    # calc avg turns made by using avg branching - avg number of links going in an angle bigger than TURN_ANGLE
    avg_turns = (360.0 - MapProblem.TURN_ANGLE * 2.0) / (360.0 / map.GetAvgBranching())
    return 1.0 / avg_turns

def SaveGammasToFile(map, filename = 'conversions.csv'):
    open(filename,'w').write('{},{},{}'.format(calc_gamma_c(map), calc_gamma_f(map), calc_gamma_s(map)))


class MapProblem(Problem):
    """The abstract class for a formal problem.  You should subclass this and
    implement the method successor, and possibly __init__, goal_test, and
    path_cost. Then you will create instances of your subclass and solve them
    with the various search functions."""

    TURN_ANGLE = 70.0

    def __init__(self, initial_junction_key, goal_junction_key, map, cost_func = None):
        self.initial = initial_junction_key
        self.goal = goal_junction_key
        self.map = map
        self.cost_func = partial(cost_func, self) if cost_func else lambda node1, link, junction2: 1
        
    def successor(self, node):
        return ((link, link.target) for link in self.map.GetJunction(node.state).links)
    
    def goal_test(self, node):
        return node.state == self.goal
    
    def path_cost(self, c, node1, link, junction2_key):
        """ Get's node1 and not junction1 because we sometimes need the path leading
        to junction1.. i.e. the C_s() func.."""
        return c + self.cost_func(node1, link, junction2_key)

if __name__ == "__main__":
    SaveGammasToFile(CountryMap().LoadMap())